package main

import (
	"context"
	"fmt"
	"image"
	"log"
	"net"
	"os"
	"sort"
	"sync"

	"github.com/nfnt/resize"
	ort "github.com/yalue/onnxruntime_go"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	_ "image/jpeg" // 支持JPEG格式
	"micro_service/proto"
)

// 定义默认的 YOLO 类别标签
var defaultYoloClasses = []string{
	"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
	"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse",
	"sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie",
	"suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
	"skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon",
	"bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut",
	"cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
	"remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book",
	"clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush",
}

// 内部边界框结构体
type boundingBox struct {
	label          string
	confidence     float32
	x1, y1, x2, y2 float32
}

// ModelSession 封装 ONNX 会话和张量

type ModelSession struct {
	Session *ort.AdvancedSession
	Input   *ort.Tensor[float32]
	Output  *ort.Tensor[float32]
}

// 为每个请求创建独立的模型会话资源池
type ModelSessionPool struct {
	mutex       sync.Mutex
	initialized bool
	sessions    chan *ModelSession
	size        int
}

// 从池中获取会话
func (p *ModelSessionPool) Acquire() (*ModelSession, error) {
	// 双重检查锁定模式确保环境只初始化一次
	if !p.initialized {
		p.mutex.Lock()
		if !p.initialized {
			// 这种情况理论上不会发生，因为我们在创建池时就初始化了环境
			p.mutex.Unlock()
			return nil, fmt.Errorf("ORT environment not initialized")
		}
		p.mutex.Unlock()
	}

	// 从通道中获取会话，如果没有可用会话则阻塞
	session := <-p.sessions
	return session, nil
}

// 将会话归还池中
func (p *ModelSessionPool) Release(session *ModelSession) {
	if session != nil {
		p.sessions <- session
	}
}

// 初始化会话池
func initSessionPool(modelPath string, poolSize int) (*ModelSessionPool, error) {
	pool := &ModelSessionPool{
		sessions: make(chan *ModelSession, poolSize),
		size:     poolSize,
	}

	// 初始化 ONNX Runtime 环境（只初始化一次）
	if err := ort.InitializeEnvironment(); err != nil {
		return nil, fmt.Errorf("error initializing ORT environment: %w", err)
	}
	pool.initialized = true

	// 预先创建指定数量的会话
	for i := 0; i < poolSize; i++ {
		session, err := createSession(modelPath, "")
		if err != nil {
			// 清理已创建的会话
			for j := 0; j < i; j++ {
				<-pool.sessions
			}
			return nil, err
		}
		pool.sessions <- session
	}

	return pool, nil
}

func createSession(modelPath, onnxruntimePath string) (*ModelSession, error) {
	// 设置 ONNX Runtime 库路径
	if onnxruntimePath != "" {
		ort.SetSharedLibraryPath(onnxruntimePath)
	}

	// 创建输入张量
	inputShape := ort.NewShape(1, 3, 640, 640)
	inputTensor, err := ort.NewEmptyTensor[float32](inputShape)
	if err != nil {
		return nil, fmt.Errorf("error creating input tensor: %w", err)
	}

	// 创建输出张量
	outputShape := ort.NewShape(1, 84, 8400)
	outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
	if err != nil {
		inputTensor.Destroy()
		return nil, fmt.Errorf("error creating output tensor: %w", err)
	}

	// 创建会话选项
	options, err := ort.NewSessionOptions()
	if err != nil {
		inputTensor.Destroy()
		outputTensor.Destroy()
		return nil, fmt.Errorf("error creating ORT session options: %w", err)
	}
	defer options.Destroy()

	// 创建会话
	session, err := ort.NewAdvancedSession(modelPath,
		[]string{"images"}, []string{"output0"},
		[]ort.ArbitraryTensor{inputTensor},
		[]ort.ArbitraryTensor{outputTensor},
		options)
	if err != nil {
		inputTensor.Destroy()
		outputTensor.Destroy()
		return nil, fmt.Errorf("error creating ORT session: %w", err)
	}

	return &ModelSession{
		Session: session,
		Input:   inputTensor,
		Output:  outputTensor,
	}, nil
}

// 实现服务接口
type yoloServiceServer struct {
	proto.UnimplementedYoloServiceServer
}

// loadImageFile 加载图像文件
// 在 loadImageFile 函数中添加
func loadImageFile(filePath string) (image.Image, error) {
	fmt.Printf("Attempting to load image from: %s\n", filePath)

	// 检查文件是否存在
	info, err := os.Stat(filePath)
	if err != nil {
		return nil, fmt.Errorf("file does not exist: %w", err)
	}

	// 检查文件大小
	if info.Size() == 0 {
		return nil, fmt.Errorf("file is empty")
	}

	fmt.Printf("File exists, size: %d bytes\n", info.Size())

	// 尝试打开文件
	file, err := os.Open(filePath)
	if err != nil {
		return nil, fmt.Errorf("error opening %s: %w", filePath, err)
	}
	defer file.Close()

	// 尝试解码文件
	img, format, err := image.Decode(file)
	if err != nil {
		return nil, fmt.Errorf("error decoding %s (detected format: %s): %w", filePath, format, err)
	}

	fmt.Printf("Successfully loaded image, format: %s\n", format)
	return img, nil
}

// prepareInput 准备模型输入
func prepareInput(img image.Image, dst *ort.Tensor[float32]) error {
	data := dst.GetData()
	channelSize := 640 * 640
	if len(data) < (channelSize * 3) {
		return fmt.Errorf("destination tensor only holds %d floats, needs %d", len(data), channelSize*3)
	}
	redChannel := data[0:channelSize]
	greenChannel := data[channelSize : channelSize*2]
	blueChannel := data[channelSize*2 : channelSize*3]

	// 调整图像大小为 640x640
	img = resize.Resize(640, 640, img, resize.Lanczos3)
	i := 0
	for y := 0; y < 640; y++ {
		for x := 0; x < 640; x++ {
			r, g, b, _ := img.At(x, y).RGBA()
			redChannel[i] = float32(r>>8) / 255.0
			greenChannel[i] = float32(g>>8) / 255.0
			blueChannel[i] = float32(b>>8) / 255.0
			i++
		}
	}

	return nil
}

// Destroy 释放会话资源
func (m *ModelSession) Destroy() {
	if err := m.Session.Destroy(); err != nil {
		log.Printf("error destroying session: %s", err)
	}
	if err := m.Input.Destroy(); err != nil {
		log.Printf("error destroying input: %s", err)
	}
	if err := m.Output.Destroy(); err != nil {
		log.Printf("error destroying output: %s", err)
	}
}

// processOutput 处理模型输出
func processOutput(output []float32, originalWidth, originalHeight int, classLabels []string) []boundingBox {
	boundingBoxes := make([]boundingBox, 0, 8400)

	var classID int
	var probability float32

	// 遍历输出数组
	for idx := 0; idx < 8400; idx++ {
		// 找到概率最高的类别
		probability = -1e9
		for col := 0; col < 80; col++ {
			if col >= len(classLabels) {
				break
			}
			currentProb := output[8400*(col+4)+idx]
			if currentProb > probability {
				probability = currentProb
				classID = col
			}
		}

		// 概率低于 0.5 的边界框被过滤掉
		if probability < 0.5 {
			continue
		}

		// 提取边界框坐标
		xc, yc := output[idx], output[8400+idx]
		w, h := output[2*8400+idx], output[3*8400+idx]
		x1 := (xc - w/2) / 640 * float32(originalWidth)
		y1 := (yc - h/2) / 640 * float32(originalHeight)
		x2 := (xc + w/2) / 640 * float32(originalWidth)
		y2 := (yc + h/2) / 640 * float32(originalHeight)

		// 添加边界框
		boundingBoxes = append(boundingBoxes, boundingBox{
			label:      classLabels[classID],
			confidence: probability,
			x1:         x1,
			y1:         y1,
			x2:         x2,
			y2:         y2,
		})
	}

	// 按置信度排序边界框
	sort.Slice(boundingBoxes, func(i, j int) bool {
		return boundingBoxes[i].confidence > boundingBoxes[j].confidence
	})

	// 非极大值抑制
	mergedResults := make([]boundingBox, 0, len(boundingBoxes))
	for _, candidateBox := range boundingBoxes {
		overlapsExistingBox := false
		for _, existingBox := range mergedResults {
			if iou(&candidateBox, &existingBox) > 0.7 {
				overlapsExistingBox = true
				break
			}
		}
		if !overlapsExistingBox {
			mergedResults = append(mergedResults, candidateBox)
		}
	}

	return mergedResults
}

// iou 计算两个边界框的交并比
func iou(a, b *boundingBox) float32 {
	// 计算交集区域
	x1 := max(a.x1, b.x1)
	y1 := max(a.y1, b.y1)
	x2 := min(a.x2, b.x2)
	y2 := min(a.y2, b.y2)

	// 如果没有交集，返回 0
	if x1 >= x2 || y1 >= y2 {
		return 0
	}

	// 计算交集面积
	intersectionArea := (x2 - x1) * (y2 - y1)

	// 计算每个边界框的面积
	aArea := (a.x2 - a.x1) * (a.y2 - a.y1)
	bArea := (b.x2 - b.x1) * (b.y2 - b.y1)

	// 计算并集面积
	unionArea := aArea + bArea - intersectionArea

	// 计算交并比
	return intersectionArea / unionArea
}

// max 返回两个 float32 中的较大值
func max(a, b float32) float32 {
	if a > b {
		return a
	}
	return b
}

// min 返回两个 float32 中的较小值
func min(a, b float32) float32 {
	if a < b {
		return a
	}
	return b
}

func (b *boundingBox) String() string {
	return fmt.Sprintf("Object %s (confidence %f): (%f, %f), (%f, %f)",
		b.label, b.confidence, b.x1, b.y1, b.x2, b.y2)
}

// 定义全局变量存储模型会话
var sessionPool *ModelSessionPool

func main() {
	// 服务器启动时初始化模型会话
	var err error
	modelPath := "./yolo11n.onnx"
	sessionPool, err = initSessionPool(modelPath, 4)
	if err != nil {
		log.Fatalf("Failed to initialize model session: %v", err)
	}
	defer func() {
		// 清理所有会话资源
		close(sessionPool.sessions)
		for session := range sessionPool.sessions {
			session.Destroy()
		}
	}()
	// 其余的服务器启动代码保持不变
	lis, err := net.Listen("tcp", ":50051")
	if err != nil {
		log.Fatalf("Failed to listen: %v", err)
	}
	s := grpc.NewServer()
	proto.RegisterYoloServiceServer(s, &yoloServiceServer{})
	fmt.Println("YOLO gRPC server listening on port 50051...")
	if err := s.Serve(lis); err != nil {
		log.Fatalf("Failed to serve: %v", err)
	}
}

// 移除重复的Detect方法实现和其他不一致的代码
func (s *yoloServiceServer) Detect(ctx context.Context, req *proto.DetectRequest) (*proto.DetectResponse, error) {
	fmt.Printf("Received detection request for image: %s\n", req.ImagePath)

	// 验证输入参数
	if req.ImagePath == "" {
		return &proto.DetectResponse{Success: false, Message: "Image path is required"},
			status.Errorf(codes.InvalidArgument, "Image path is required")
	}

	// 从池中获取会话
	session, err := sessionPool.Acquire()
	if err != nil {
		return &proto.DetectResponse{Success: false, Message: "Failed to acquire model session"},
			status.Errorf(codes.Internal, "Failed to acquire model session: %s", err)
	}
	defer sessionPool.Release(session)

	// Read the input image into an image.Image object
	pic, e := loadImageFile(req.ImagePath)
	if e != nil {
		fmt.Printf("error loading input image: %s\n", e)
		return nil, status.Errorf(codes.Internal, "error loading input image: %s", e)
	}

	originalWidth := pic.Bounds().Canon().Dx()
	originalHeight := pic.Bounds().Canon().Dy()

	e = prepareInput(pic, session.Input)
	if e != nil {
		fmt.Printf("error converting image to network input: %s\n", e)
		return nil, status.Errorf(codes.Internal, "error converting image to network input: %s", e)
	}

	e = session.Session.Run()
	if e != nil {
		fmt.Printf("error running ORT session: %s\n", e)
		return nil, status.Errorf(codes.Internal, "error running ORT session: %s", e)
	}
	classLabels := defaultYoloClasses
	if len(req.ClassLabels) > 0 {
		classLabels = req.ClassLabels
	}
	// Print the results
	boxes := processOutput(session.Output.GetData(), originalWidth, originalHeight, classLabels)
	for i, box := range boxes {
		fmt.Printf("Box %d: %s\n", i, &box)
	}
	// 转换为 gRPC 响应格式
	responseBoxes := make([]*proto.BoundingBox, 0, len(boxes))
	for _, box := range boxes {
		responseBoxes = append(responseBoxes, &proto.BoundingBox{
			X1:         box.x1,
			Y1:         box.y1,
			X2:         box.x2,
			Y2:         box.y2,
			Label:      box.label,
			Confidence: box.confidence,
		})
	}

	return &proto.DetectResponse{
		Boxes:   responseBoxes,
		Success: true,
		Message: "Detection completed successfully",
	}, nil
}
